import torch
from utils import *
import numpy as np

def test(soft, test_loader, model):
#     model.eval()
    sdr_list_s = []
#     sdr_list_n = []
    sdr_list_rs = []
    sdr_list_rx = []
    sdr_list_diff = []
    max_rs_sdr = 0
    max_rx_sdr = 0
    max_rs = None
    max_rx = None

    i = 0
    for wave_s, wave_x, source, mixture in test_loader: 

        source = source[0]
        mixture = mixture[0]
#         print(data)
        noise = (mixture-source).cuda()
        source = source.cuda()
        mixture = mixture.cuda()
        s_h, n_h, arg_s, arg_n = model(mixture, soft=soft)
        
        sdr_list_s.append(SISDR(source, s_h))    
#         sdr_list_n.append(SISDR(noise, n_h))
        
        rs = rebuild(s_h).unsqueeze(dim = 0)
        rn = rebuild(n_h).unsqueeze(dim = 0)
        
        rs_sdr = SISDR(wave_s, rs)
        rx_sdr = SISDR(wave_x, rs+rn)
        sdr_list_rs.append(rs_sdr)
        sdr_list_rx.append(rx_sdr)
        
        start_sdr = SISDR(wave_s, wave_x/max(wave_x[0]))
        sdr_list_diff.append(rs_sdr-start_sdr)
        
        if max_rs_sdr < rs_sdr:
            max_rs_sdr = rs_sdr
            max_rs = (wave_x, rs, rn)
#         if max_rx_sdr < rx_sdr:
#             max_rx_sdr = rx_sdr
#             max_rx = (rs, rn)
         
    
    s_score = np.mean(sdr_list_s) 
#     n_score = np.mean(sdr_list_n)
    rs_score = np.mean(sdr_list_rs)
    rx_score = np.mean(sdr_list_rx)
    diff_score = np.mean(sdr_list_diff)
    
    return s_score, rs_score, rx_score, diff_score, max_rs, arg_s, arg_n